"""
通用工具函数
提供各种辅助功能
"""
import random
import numpy as np
import torch
import os
import time
from pathlib import Path


def set_seed(seed=42):
    """
    设置随机种子以确保可复现性
    
    参数:
        seed: 随机种子值
    """
    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    
    if torch.cuda.is_available():
        torch.cuda.manual_seed(seed)
        torch.cuda.manual_seed_all(seed)
        torch.backends.cudnn.deterministic = True
        torch.backends.cudnn.benchmark = False
    
    # 设置Python hash seed
    os.environ['PYTHONHASHSEED'] = str(seed)
    
    print(f"Random seed set to {seed}")


def check_requirements(requirements=None):
    """
    检查依赖包是否安装
    
    参数:
        requirements: 依赖包列表，如果为None则使用默认列表
    
    返回:
        bool: 所有依赖是否都已安装
    """
    if requirements is None:
        requirements = [
            'torch', 'numpy', 'pandas', 'pillow', 
            'tqdm', 'pyyaml', 'scikit-learn', 'scipy'
        ]
    
    missing = []
    
    for pkg in requirements:
        try:
            __import__(pkg)
        except ImportError:
            missing.append(pkg)
    
    if missing:
        print(f"Missing packages: {', '.join(missing)}")
        print(f"Install with: pip install {' '.join(missing)}")
        return False
    
    print("✓ All required packages are installed")
    return True


def colorstr(*args):
    """
    为字符串添加颜色（终端显示）
    
    参数:
        *args: 'color', 'string'
    
    返回:
        colored_string: 带颜色的字符串
    
    示例:
        >>> print(colorstr('blue', 'bold', 'hello world'))
    """
    colors = {
        'black': '\033[30m',
        'red': '\033[31m',
        'green': '\033[32m',
        'yellow': '\033[33m',
        'blue': '\033[34m',
        'magenta': '\033[35m',
        'cyan': '\033[36m',
        'white': '\033[37m',
        'bright_black': '\033[90m',
        'bright_red': '\033[91m',
        'bright_green': '\033[92m',
        'bright_yellow': '\033[93m',
        'bright_blue': '\033[94m',
        'bright_magenta': '\033[95m',
        'bright_cyan': '\033[96m',
        'bright_white': '\033[97m',
        'end': '\033[0m',
        'bold': '\033[1m',
        'underline': '\033[4m'
    }
    
    string = args[-1]
    color_codes = ''.join(colors.get(x, '') for x in args[:-1])
    return f'{color_codes}{string}{colors["end"]}'


def increment_path(path, exist_ok=False, sep='', mkdir=False):
    """
    自动递增路径，例如 runs/exp --> runs/exp{sep}2, runs/exp{sep}3, ...
    
    参数:
        path: 基础路径
        exist_ok: 如果为True，则返回原路径
        sep: 分隔符
        mkdir: 是否创建目录
    
    返回:
        path: 递增后的路径
    
    示例:
        >>> path = increment_path('runs/exp')  # runs/exp
        >>> path = increment_path('runs/exp')  # runs/exp2
    """
    path = Path(path)
    
    if path.exists() and not exist_ok:
        path, suffix = (path.with_suffix(''), path.suffix) if path.is_file() else (path, '')
        
        # 查找已存在的编号
        for n in range(2, 9999):
            p = f'{path}{sep}{n}{suffix}'
            if not Path(p).exists():
                path = Path(p)
                break
    
    if mkdir:
        path.mkdir(parents=True, exist_ok=True)
    
    return path


def make_divisible(x, divisor):
    """
    使数字可被divisor整除
    
    参数:
        x: 原始数字
        divisor: 除数
    
    返回:
        int: 调整后的数字
    """
    return int(np.ceil(x / divisor) * divisor)


def check_img_size(img_size, stride=32):
    """
    检查图像尺寸，确保可被stride整除
    
    参数:
        img_size: 图像尺寸 (int 或 tuple)
        stride: 步长
    
    返回:
        img_size: 调整后的图像尺寸
    """
    if isinstance(img_size, int):
        new_size = make_divisible(img_size, stride)
        if new_size != img_size:
            print(f'Warning: image size {img_size} adjusted to {new_size} (divisible by {stride})')
        return new_size
    else:
        new_size = [make_divisible(x, stride) for x in img_size]
        if new_size != list(img_size):
            print(f'Warning: image size {img_size} adjusted to {tuple(new_size)} (divisible by {stride})')
        return tuple(new_size)


def count_parameters(model, trainable_only=False):
    """
    计算模型参数数量
    
    参数:
        model: 模型
        trainable_only: 是否只计算可训练参数
    
    返回:
        count: 参数数量
    """
    if trainable_only:
        count = sum(p.numel() for p in model.parameters() if p.requires_grad)
    else:
        count = sum(p.numel() for p in model.parameters())
    
    return count


def print_model_info(model, input_size=None):
    """
    打印模型信息
    
    参数:
        model: 模型
        input_size: 输入尺寸（可选）
    """
    print("\n" + "="*70)
    print("Model Information".center(70))
    print("="*70)
    
    # 参数数量
    total_params = count_parameters(model, trainable_only=False)
    trainable_params = count_parameters(model, trainable_only=True)
    
    print(f"\nParameters:")
    print(f"  Total: {total_params:,}")
    print(f"  Trainable: {trainable_params:,}")
    print(f"  Non-trainable: {total_params - trainable_params:,}")
    
    # 模型大小
    param_size = sum(p.numel() * p.element_size() for p in model.parameters())
    buffer_size = sum(b.numel() * b.element_size() for b in model.buffers())
    size_mb = (param_size + buffer_size) / 1024 / 1024
    
    print(f"\nModel Size:")
    print(f"  Memory: {size_mb:.2f} MB")
    
    # 测试前向传播（如果提供了input_size）
    if input_size is not None:
        try:
            device = next(model.parameters()).device
            dummy_input = torch.randn(input_size).to(device)
            model.eval()
            with torch.no_grad():
                output = model(dummy_input)
            print(f"\nForward Pass Test:")
            print(f"  Input shape: {input_size}")
            print(f"  Output shape: {output.shape}")
        except Exception as e:
            print(f"\nForward Pass Test: Failed - {e}")
    
    print("="*70 + "\n")


def time_sync():
    """PyTorch精确时间同步"""
    if torch.cuda.is_available():
        torch.cuda.synchronize()
    return time.time()


def select_device(device=''):
    """
    选择设备 - 修复版本，正确处理默认参数
    
    参数:
        device: 设备字符串
            - '': 空字符串，自动选择（有GPU用GPU，无GPU用CPU）
            - 'cpu': 强制使用CPU
            - 'cuda': 使用默认GPU
            - '0', '1', 'cuda:0': 指定GPU编号
    
    返回:
        device: torch.device对象
    """
    # 处理字符串设备
    device_str = str(device).lower().strip()
    cpu = device_str == 'cpu'
    
    if cpu:
        # 强制使用CPU
        os.environ['CUDA_VISIBLE_DEVICES'] = '-1'
    elif device_str and device_str != 'cuda':
        # 指定GPU编号（如 '0', '1', 'cuda:0'）
        gpu_id = device_str.replace('cuda:', '')
        os.environ['CUDA_VISIBLE_DEVICES'] = gpu_id
    
    # 检查CUDA是否可用
    cuda = not cpu and torch.cuda.is_available()
    device_obj = torch.device('cuda:0' if cuda else 'cpu')
    
    # 打印设备信息
    if cuda:
        print(f'Using device: {device_obj}')
        print(f'  GPU: {torch.cuda.get_device_name(0)}')
        print(f'  Memory: {torch.cuda.get_device_properties(0).total_memory / 1024**3:.2f} GB')
    else:
        print(f'Using device: cpu')
        if not cpu and device_str:
            print(f'  Note: CUDA requested but not available')
    
    return device_obj


if __name__ == "__main__":
    # 测试工具函数
    print("Testing general utilities...")
    
    # 测试设置随机种子
    print("\n1. Testing set_seed...")
    set_seed(42)
    
    # 测试检查依赖
    print("\n2. Testing check_requirements...")
    check_requirements()
    
    # 测试颜色字符串
    print("\n3. Testing colorstr...")
    print(colorstr('green', 'bold', 'Success!'))
    print(colorstr('red', 'Error!'))
    print(colorstr('blue', 'Info:'), 'This is information')
    
    # 测试路径递增
    print("\n4. Testing increment_path...")
    import tempfile
    temp_dir = tempfile.mkdtemp()
    path1 = increment_path(Path(temp_dir) / 'exp', mkdir=True)
    path2 = increment_path(Path(temp_dir) / 'exp', mkdir=True)
    print(f"  Path 1: {path1}")
    print(f"  Path 2: {path2}")
    
    # 测试参数计数
    print("\n5. Testing count_parameters...")
    import torch.nn as nn
    model = nn.Sequential(
        nn.Linear(10, 20),
        nn.ReLU(),
        nn.Linear(20, 5)
    )
    print(f"  Total parameters: {count_parameters(model):,}")
    
    # 测试设备选择
    print("\n6. Testing select_device...")
    device = select_device('')
    
    print("\n✓ General utilities test passed!")